#!/usr/bin/python
# -*- coding: utf-8 -*-
# @version        : 1.0
# @Create Time    : 2024/05/07 17:23
# @File           : hpc_api.py
# @IDE            : PyCharm
# @desc           : 调用hpc_api

"""
aiohttp 异步请求文档：https://docs.aiohttp.org/en/stable/client_quickstart.html
"""
from enum import Enum
import time
from typing import List, Optional, Union
from aiohttp import ClientError, ClientResponseError
import aiohttp
import jsonschema
from pydantic import BaseModel, ConfigDict, field_validator
from utils.tools import filter_config, get_nested_dict, render_config
from core.logger import logger
from application import settings
from core.database import mongo_db
from apps.base.record.crud import ApiCallRecordDal


class ApiType(Enum):
    DETAIL: str = "detail"
    LIST: str = "list"
    PUSH: str = "push"


class ApiRequest(BaseModel):
    url: str
    method: Optional[str] = "GET"
    headers: Optional[dict] = {}
    json: Optional[Union[dict, str, list]] = None
    params: Optional[Union[dict, str, list]] = None
    data: Optional[Union[dict, str, list]] = None


class ApiResponse(BaseModel):
    status_code: Optional[Union[List[int], int, str]] = [200]
    type: Optional[str] = "json"
    fields: Optional[dict] = {}  # 字段映射
    schema: Optional[dict] = {}  # 字段校验，参考json_schema

    @field_validator('status_code')
    def fomat_status_code(cls, v):
        if isinstance(v, str):
            if v.isdigit():
                v = int(v)
            else:
                raise ValueError(f'status code {v} unvalid')
        elif isinstance(v, list):
            return v
        return [v]


class ApiConfig(BaseModel):
    model_config = ConfigDict(from_attributes=True)
    name: Optional[str] = "未命名"
    ssl: Optional[bool] = False

    env: Optional[dict] = {}
    request: Optional[ApiRequest]
    response: Optional[ApiResponse]


class ApiClient:
    def __init__(self, config: dict, headers = None, pyload=None, **kwargs):
        self.config = ApiConfig(**config)
        if headers: self.config.request.headers.update(headers)
        if not self.config.env:
            self.config.env = {}
        self.config.env.update(**kwargs)
        if pyload:
            if self.config.request.json == "pyload":
                self.config.request.json = pyload
            if self.config.request.data == "pyload":
                self.config.request.data = pyload
            if self.config.request.params == "pyload":
                self.config.request.params = pyload

    @staticmethod
    def _validate_resposne(response_data, schema):            
        try:
            jsonschema.validate(instance=response_data, schema=schema)
        except jsonschema.exception.ValidattionError as e:
            raise ValueError(f"返回数据检验不匹配：{schema}")

    def _map_fields(self, response_data, fields):
        if not fields:
            return response_data
        
        mapped_data = {}
        for field, field_config in fields.items():
            if isinstance(field_config, str):
                mapped_data[field] = get_nested_dict(response_data, field_config)
            elif isinstance(field_config, dict):
                path = field_config.get('path', field)
                value = get_nested_dict(response_data, path)
                
                if 'fields' in field_config:
                    if isinstance(value, list):
                        # 处理列表中每个对象的字段映射
                        mapped_data[field] = [self._map_fields(v, field_config['fields']) for v in value]
                    elif isinstance(value, dict):
                        mapped_data[field] = self._map_fields(value, field_config['fields'])
                    else:
                        logger.warning(f"数据类型{type(value)}无法经配置进行映射：{field_config['fields']}")
                else:
                    mapped_data[field] = value
        return mapped_data

    async def _handle_response(self, resp):
        response = self.config.response
        if response.status_code:
            if resp.status not in response.status_code:
                msg = await resp.text()
                logger.error(f"接口[{self.config.name}]期待返回为 {response.status_code}, 实际返回为 {resp.status}, 返回文本信息为：{msg}")
                raise RuntimeError(f"接口[{self.config.name}]期待返回为 {response.status_code}, 实际返回为 {resp.status}")

        if response.type.lower() == 'json':
            data = await resp.json()
        elif response.type.lower() == 'text':
            data = await resp.text()
        else:
            logger.error(f"接口[{self.config.name}]使用了不支持的返回类型解析结果，{response.type}")
            raise RuntimeError(f"接口[{self.config.name}]配置了不支持的返回类型：{response.type}")

        logger.debug(f"接口返回数据为：{data}")

        if response.schema:
            self._validate_resposne(data, response.schema)

        return self._map_fields(data, response.fields)

    def apicall_record_wrapper(func):
        async def wrapper(self, params):
            start_time = time.time()
            error = None
            try:
                resp = func(self, params)
            except Exception as e:
                error = e
                resp = None
                raise
            finally:
                document = {
                    'process_time': time.time() - start_time,
                    'error': error,
                    'response_body': None,
                    'params': params,
                    'config': self.config.model_dump()
                }
                if settings.MONGO_DB_ENABLE and settings.APICALL_LOG_RECORD \
                    and self.config.request.method.upper() in settings.APICALL_RECORD_METHOD \
                    and self.config.name not in settings.IGNORE_APICALL_FUNCTION:
                    
                    # async for mongo in mongo_db:
                    if resp:
                        async for res in resp:
                            document['response_body'] = await res.text()
                            document['response_status'] = res.status
                            await ApiCallRecordDal(mongo_db).create_data(document)
                            yield res
                    else:
                        await ApiCallRecordDal(mongo_db).create_data(document)
                        async for res in resp: yield res
                else:
                    if resp:
                        async for res in resp: yield res
                    else:
                        yield resp     
        return wrapper

    @apicall_record_wrapper
    async def _requet(self, params):
        async with aiohttp.ClientSession(connector=aiohttp.TCPConnector(ssl=self.config.ssl)) as session:
            async with session.request(
                **params
            ) as resp:
                yield resp
            
    async def request(self):
        params = render_config(self.config.request.model_dump(), self.config.env)
        params = filter_config(params)
        logger.debug(f"请求外部接口，参数为：{params}")
        try:
            async for resp in self._requet(params=params):
                logger.debug(f"请求外部接口 [{self.config.name}]，响应为：{resp}")
                if resp: return await self._handle_response(resp)
        except ClientResponseError as e:
            raise RuntimeError(f'客户端响应错误: {e.status}') from e
        except ClientError as e:
            raise RuntimeError(f'请求客户端错误: {e}') from e
        except RuntimeError as e:
            raise
        except Exception as e:
            raise RuntimeError(f"未知错误: {e}") from e


class ListApiClient(ApiClient):
    def __init__(self, config, **kwargs):
        super().__init__(config, **kwargs)

        self.page = 1
        self.limit = kwargs.get("limit") or self.config.env.get("limit") or self.config.env.get("page_size") or 20
        self.total_count = 0

    async def request(self):
        self.config.env.update(
            {
                "page": self.page,
                "limit": self.limit,
                "page_size": self.limit
            }
        )
        return await super().request()

    async def page_caller(self):
        resp = await self.request()
        self.total_count = resp.get("total_count")
        if self.total_count is None:
            raise RuntimeError(f"分页查询未获取到总记录数，total_count: {self.total_count}")
        if isinstance(self.total_count, str) and not self.total_count.isdigit():
            raise RuntimeError(f"分页查询获取到总记录数类型有误，total_count: {self.total_count}")
        
        logger.debug(f"分页查询，获取到总记录数 total_count: {self.total_count}")
        logger.debug(f"页码 {self.page} 响应: {resp}")
        yield resp

        while self.page * self.limit < int(self.total_count):
            self.page += 1
            resp = await self.request()
            logger.debug(f"页码 {self.page} 响应: {resp}")
            yield resp
